open Filename
open Printf
open Ext
open Str
open Circuit 
open Unix 
open Random
open Bn
(* Globals used for command line parameters *)

let network = ref ""
let acfile = ref ""
let qfile = ref ""
let vfile = ref ""
let evfile = ref ""
let datafile = ref ""
let l1param = ref 0.0
let opt_l1param = ref 0.0
let stddev = ref 10.0
let prior = ref 0.0001
let edge_penalty = ref 10.0
let num_comps = ref 4
let lambda = ref 1.0
let scale = ref 1.0
let schemafile = ref ""
let outfile = ref ""
let outmodel = ref "" 
let output_dir = ref ""
let seed = ref 0 
(*((int_of_float Unix.time() )  mod  1000) 
*)
let override = ref false

let load_circuit filename =
  try Circuit.load_with_features (open_in filename) 
  with Circuit.Unknown_child(n,c) -> 
         nlogf "Error: node %s references unknown child node %s.\n" n c; 
         exit 0 

let usage = "Usage: mtlearn -i <data> -s <schema> -o <directory-model> [...]"
let args = Arg.align
  ([
    (*("-va", Arg.Set_string vfile, " Validatione file");*)
    ("-i", Arg.Set_string datafile, " Training data file");
    (*("-q",  Arg.Set_string qfile,  " Query file"); *)

    ("-s",  Arg.Set_string schemafile,  " Data schema (optional)");
  (*("-of", Arg.Set_string outfile, " Output inference file");*)
    ("-o", Arg.Set_string output_dir, " Output model directory");

    
    ("-k",  Arg.Set_int num_comps,  " Number of component");
    ("-seed",  Arg.Set_int seed,  " Random seed");
    ("-f",  Arg.Set override,  " Force to override model directory");
    (*("-net", Arg.Set_string network, " Network name") *)
  ]
    @ common_arguments)

let load_schema file =
  let comma = regexp "[ \t]*,[ \t]*" in
  let channel = open_in file in 
  let line = input_line channel in 
  let tokens = split comma line in
  let schema = List.map (function s -> int_of_string s) tokens in
  schema



let load_data file = 
  let data = ref [] in 
  let channel = open_in file in 
  let comma = regexp "[ \t]*,[ \t]*" in
  try
    while true; do
      let line = input_line channel in 
      let tokens = split comma line in
      let sample = Array.map (function s -> int_of_string s) (Array.of_list tokens) in
      (*List.iter (fun s-> print_int s; print_string " ") sample;
      *)
      (*data := List.append !data [sample]; *)
      data := sample :: !data;
    done;
    let d = List.rev !data in 
    d
  with End_of_file -> close_in_noerr channel;
    let d = List.rev !data in
    d
    (*!data*)

let absf f = if f < 0.0 then -.f else f

let distance x y =
  let xf = Array.map float_of_int x in
  let fsquare a = a *. a in
  let f = Array.map2 ( fun a b -> fsquare ( a -. b ) ) xf y in
  let d = Array.sumf f in
  d /. 2.0

let log2 x = log x /. log 2.0 



let print_sample out sample =
  let line = String.concat "," (List.map string_of_int sample) in 
  Printf.fprintf out "%s\n" line

let save_data data filename = 
  let out = open_out filename in  
  List.iter (print_sample out) data

let save_data_parts data_lists num_comps =
  for i=0 to num_comps - 1 do
    let filename = Printf.sprintf "%s.ts.%d.data" !network (i+1) in
    let data_part  = data_lists.(i) in
    save_data data_part filename  
  done

let rem_first l = 
  match l with            
    | [] -> []
  | h::t -> t 
   

let soft_counts schema data comp_probs k=
  (* Collect all pairwise counts *)
  let numvars = Array.length schema in
  let marg_counts = Array.map (fun d -> Array.create d 0.0) schema in
  let create_counts i j = Array.create_matrix schema.(i) schema.(j) 0.0 in
  let joint_counts = 
    Array.init numvars (fun i -> Array.init i (create_counts i)) in
  let add_counts x p=
    for i = 0 to numvars - 1 do
      let xi = x.(i) in
      marg_counts.(i).(xi) <- marg_counts.(i).(xi) +. exp p.(k);
      for j = 0 to i - 1 do
        let xj = x.(j) in
        joint_counts.(i).(j).(xi).(xj) <- joint_counts.(i).(j).(xi).(xj) +. exp p.(k)
      done;
    done in
  List.iter2 (add_counts) data comp_probs;
  (marg_counts, joint_counts, List.length data)

let compute_mi schema num_examples marg_counts joint_counts comp_probs k=
  (* Compute a single mutual information score *)
  let tp = List.fold_left (fun s p -> exp p.(k) +. s ) 0.0 comp_probs in
  let total = !prior +. tp in
  (*
 let total = !prior +. float_of_int num_examples in
  *)
  let calc_mi i j =
    let mi = ref 0.0 in
    let ip = !prior /. (float_of_int schema.(i)) in
    let jp = !prior /. (float_of_int schema.(j)) in
    let ijp = !prior /. (float_of_int (schema.(i) * schema.(j))) in
    for ival = 0 to schema.(i) - 1 do
      for jval = 0 to schema.(j) - 1 do
        let p_ij = 
          (ijp +. joint_counts.(i).(j).(ival).(jval)) /. total in
        let p_i  = (ip +. marg_counts.(i).(ival)) /. total in
        let p_j  = (jp +. marg_counts.(j).(jval)) /. total in
        dlogf "P(x_%d = %d, x_%d = %d) = %f\n" i ival j jval p_ij; 
        if p_ij > 0. then
          mi := !mi +. p_ij *. log (p_ij /. (p_i *. p_j))
      done;
    done;
    !mi in

  (* Calculate all mutual informations *)
  let numvars = Array.length schema in
  let all_mi = Array.init numvars (fun i -> Array.init i (calc_mi i)) in
  all_mi


let rec add_edge all_mi visited edges chosen_edges i j =

  let numvars = Array.length visited in

  (* Add new edge *)
  let chosen_edges = (i, j) :: chosen_edges in

  (* Add linked edges... *)
  let n = if visited.(i) then j else i in

  let get_mi n' =
    if n = n' then (n, n', 0.) 
    else begin
      let n1 = min n n' and n2 = max n n' in
      let mi = all_mi.(n2).(n1) in
      dlogf "I(%d ; %d) = %f\n" n n' mi; 
      (n, n', mi) 
    end in
  let new_edges = Array.init numvars get_mi in
  Array.iter (Heap.add edges) new_edges;
  visited.(n) <- true;

  (* Remove candidate edges that would create a cycle *)
  Heap.remove_all edges (fun (i, j, v) -> visited.(i) && visited.(j));

  (* Check to see if we're done.  (We'll end up with numvars-1 edges in
     the end, but for now, we have a dummy edge at the beginning.) *)
  if List.length chosen_edges = numvars then
    chosen_edges
  else begin
    (* Select the best edge and recurse! *)
    let (i, j, mi) = Heap.min edges in 
    dlogf "CHOSEN: %d -> %d (%f)\n" i j mi; 
    add_edge all_mi visited edges chosen_edges i j
  end

let chow_liu schema all_mi =
  let numvars = Array.length schema in
  let visited = Array.create numvars false in
  let edge_comp (_, _, mi1) (_, _, mi2) = mi1 > mi2 in
  let edges = Heap.create edge_comp 100 in
  let all_edges = add_edge all_mi visited edges [] 0 0 in
  (* Remove the first dummy edge *)
  List.tl (List.rev all_edges)

(*
let rec add_edge all_mi visited edges chosen_edges i j =
  let numvars = Array.length visited in
  (* Add new edge *)
  let chosen_edges = (i, j) :: chosen_edges in
  (* Add linked edges... *)
  let n = if visited.(i) then j else i in
  let get_mi n' =
    if n = n' then (n, n', 0.) 
    else begin
      let n1 = min n n' and n2 = max n n' in
      let mi = all_mi.(n2).(n1) in
      dlogf "I(%d ; %d) = %f\n" n n' mi; 
      (n, n', mi) 
    end in
  let new_edges = Array.init numvars get_mi in
  Array.iter (Heap.add edges) new_edges;
  visited.(n) <- true;

  (* Remove candidate edges that would create a cycle *)
  Heap.remove_all edges (fun (i, j, v) -> visited.(i) && visited.(j));
  (* Check to see if we're done.  (We'll end up with numvars-1 edges in
     the end, but for now, we have a dummy edge at the beginning.) *)
  if List.length chosen_edges = numvars then
    chosen_edges
  else begin
    (* Select the best edge and recurse! *)
    let (i, j, mi) = Heap.min edges in 
    dlogf "CHOSEN: %d -> %d (%f)\n" i j mi; 
    add_edge all_mi visited edges chosen_edges i j
  end
*)

let create_bn schema edges marg_counts joint_counts =

  (* Create a BN from the edge set and the mutual informations *)
  let bn = Bn.create_empty_network schema in

  (* i = parent; j = child *)
  let set_cpd (i, j) =
    let counts = 
      if i < j then 
        Array.transpose joint_counts.(j).(i)
      else
        joint_counts.(i).(j) in
    let table = Array.create_matrix schema.(i) schema.(j) 0. in
    let ip = !prior /. (float_of_int schema.(i)) in
    let ijp = !prior /. (float_of_int (schema.(i) * schema.(j))) in
    for ii = 0 to schema.(i) - 1 do
      for jj = 0 to schema.(j) - 1 do
        table.(ii).(jj) <- log (((ijp +.  counts.(ii).(jj))) /. 
                                (ip +.  marg_counts.(i).(ii)));
      done
    done;
    Bn.set_cpt bn j [i] table in 
  List.iter set_cpd edges; 

  (* Set marginal distributions for all nodes with no parents *)
  for i = 0 to Array.length schema - 1 do
    if Bn.parents bn i = [] then begin
      let ip = !prior /. (float_of_int schema.(i)) in
      let denom = !prior +.  (Array.sumf marg_counts.(i)) in
      let count_to_weight c = log ((ip +. c) /. denom) in
      let dist = Array.map count_to_weight marg_counts.(i) in
      Bn.set_cpt bn i [] [|dist|]
    end 
  done;
  bn



let answer_query bn d =
  Bn.loglikelihood bn d 
  (* do somthing here to compute Tk(xi)*)
  
let answer_mixture_query bn_l w_l d =
  let probs = List.map2 (fun bn logw -> let logp = answer_query bn d in logw +. logp) bn_l w_l in
  probs

let get_final_prob bn_l w_l d = 
  let comp_probs = answer_mixture_query bn_l w_l d in
  let p = logsumexp comp_probs in
  p

let learn_mixture_params schema data num_comps=
  let unnormal_ws = Array.create num_comps 0.0  in
  (*let t = Unix.time() in
  
  let seed = (int_of_float t)  mod  500 in
  Random.init seed;*)
  for k=0 to num_comps - 1 do
    unnormal_ws.(k) <- Random.float 1.0;
  
  done;
  let w_sum = Array.sumf unnormal_ws in 
  let ws = Array.map (fun w->log w -. log w_sum) unnormal_ws in
  
  let bns = ref [||] in 
  let round = ref 0 in
  let err = ref 10.0 in
  let n_f = float_of_int (List.length data) in 
  let pk_rn = ref [] in
  for i = 1 to List.length data do  

      let x = Array.create num_comps 0.0 in 
      for k=0 to num_comps - 1 do
          x.(k) <- log (Random.float 1.0); 
      done;  
      (*pk_rn := List.append !pk_rn [x];*)
      pk_rn := x :: !pk_rn
  done;
  let rev_pl = List.rev !pk_rn in
  pk_rn := rev_pl;

  nlogf "Intial CL trees\n"; flush_all();
  for k=0 to num_comps - 1 do
      let (marg_counts, joint_counts, n) = soft_counts schema data !pk_rn k in
      let mi_k = compute_mi schema n marg_counts joint_counts !pk_rn k in
      let edges_k = chow_liu schema mi_k in
      let bn_k = create_bn schema edges_k marg_counts joint_counts in
      bns := Array.append !bns [|bn_k|];
  done;

  let sgamma = ref [||] in
  let sgamma_l = ref [] in 
  for i=0 to (List.length data) - 1 do
      (*sgamma := Array.append !sgamma [|Array.create num_comps 0.0|];*)
      sgamma_l := (Array.create num_comps 0.0)::!sgamma_l
  done;
  sgamma := Array.of_list !sgamma_l;
  while !round < 100 && !err > 0.01 do
    (*Printf.nlogf "------------------ Round %d -------------------\n" !round; flush_all();
    *)
    incr round;

    let gama =  Array.create num_comps 0.0 in

    
    let update_gama d_i d = (* gama is Gamma *)
      (*nlogf "answer_mixture_query %d\n" d_i; flush_all();  *)
      let cs = answer_mixture_query (Array.to_list !bns) (Array.to_list ws) d in 
      let s = logsumexp cs in 
      (*nlogf "value %d:%f\n" d_i s; flush_all(); *)
      Array.iteri (fun k p -> !sgamma.(d_i).(k) <- (p -. s);
      gama.(k)<- log (exp gama.(k) +. exp !sgamma.(d_i).(k))(*; Printf.nlogf "%d %d %f\n" d_i k sgamma.(d_i).(k)*)) (Array.of_list cs); 
      s in


    let data_ar = Array.of_list data in
    (*nlogf "Update gama\n"; flush_all();*)
      let probs = Array.mapi (update_gama) data_ar in 
    ignore(probs);
    (*nlogf "Compute pks\n"; flush_all();*)
    let p_ks = Array.mapi (fun d_i s_gamma_i -> 
            Array.mapi (fun k v -> s_gamma_i.(k) -. gama.(k) ) s_gamma_i) !sgamma in
    let p_kl = Array.to_list p_ks in
    err := 0.0;

    (*nlogf "Create new CL trees\n"; flush_all();*)
    for k=0 to num_comps - 1 do
      let (marg_counts, joint_counts, n) = soft_counts schema data p_kl k in
      let mi_k = compute_mi schema n marg_counts joint_counts p_kl k in
      let edges_k = chow_liu schema mi_k in
      let bn_k = create_bn schema edges_k marg_counts joint_counts in
      !bns.(k) <- bn_k;
      let w_k = gama.(k) -. log n_f in
      err := !err +. absf ( w_k -. ws.(k) );
      ws.(k) <- w_k;
    done;
    nlogf "round %d, error = %f\n" !round !err; 
  done;
  (!bns, ws)
      

let remove_nonempty_dir dirname = 
  let filelist = Sys.readdir dirname in 
  Array.iter (fun f->Sys.remove (String.concat dir_sep [dirname; f]) )filelist;
  ignore()


;;
let main () = 
  network := "spac";
  Timer.start "total";

  (* Parse arguments *)
  seed := ((int_of_float (Unix.time()) )  mod  1000); 
  Arg.parse args ignore usage;
  if !datafile = "" || (!output_dir = "") then
      (Arg.usage args usage; exit 1);

  Random.init !seed;
  if not (check_suffix !output_dir ".spn") then begin nlogf "Error: output directory model should end with .spn \n"; exit 1 end;
  
  let odir = sprintf "%s" !output_dir in
	  if Sys.file_exists odir then 
      begin
        if not !override then begin
          nlogf "Error: directory %s exists\n"!output_dir; exit 1 end
        else begin
          remove_nonempty_dir !output_dir
        end
      end
    else Unix.mkdir odir 0o755;


	
  let datas = load_data !datafile in 
  let schema = 
    if !schemafile <> "" then
      load_schema !schemafile
    else Array.to_list (Data.schema datas) in 
  
  Unix.chdir odir;
	   
  if !qfile <> "" then 
    if !outfile = "" then begin 
      nlogf "Needs output inference file as well\n"; 
          (Arg.usage args usage; exit 0)
  end;
  common_log_init ();
  nlogf "Loading data....\n";
  flush_all();



  
  let t1 = Sys.time() in

  nlogf "Learning params....\n";
  flush_all();

  let (bn, ws) = learn_mixture_params (Array.of_list schema) datas !num_comps in
  let t2 = Sys.time() in
  let bn_l = Array.to_list bn in
  let w_l = Array.to_list ws in
  flush_all();
  let data_llg = (List.sumf_map ( fun d->get_final_prob bn_l w_l d ) datas) /. float_of_int (List.length datas) in

  let valid_llg = ref 0.0 in
  if ( !vfile <> "" ) then begin
    let val_data = load_data !vfile in  

    valid_llg := (List.sumf_map ( fun d->get_final_prob bn_l w_l d ) val_data) /. float_of_int (List.length val_data)
  end;
  let test_llg = ref 0.0 in 
  if !qfile <> "" && !outfile <> "" then begin
    let queries = load_data !qfile in 
    let output_result = open_out !outfile in
    test_llg := (List.sumf_map ( fun d->let logp = get_final_prob bn_l w_l d in fprintf output_result "%f\n" logp; logp  ) queries) /. float_of_int (List.length queries);
    close_out_noerr output_result
  end;
  outmodel := "spac.m";
  if !outmodel <> "" then
  begin
    for i=1 to !num_comps do
      let name = sprintf "%s-%d.bn" !network i in
      Bn.write_auto name bn.(i-1);

      let acname = sprintf "%s-%d.ac" !network i in
      let cmd = Printf.sprintf "acve -m %s -o %s" name acname in 
      nlogf "%s\n" cmd;
	  ignore(Unix.system cmd);
	  
	  ignore();
    done;

    let mout = open_out !outmodel in
    fprintf mout "%s\n%d\n" !network (!num_comps + 1);
    fprintf mout "n 0 + -1\n";

    for i=1 to !num_comps do
      fprintf mout "%d " i
    done;
    fprintf mout "\n";
    for i=0 to !num_comps - 1 do
      fprintf mout "%f " ws.(i)
    done;

    fprintf mout "\n";
    for i=0 to (List.length schema) - 1  do
      fprintf mout "%d " i;
    done;
    fprintf mout "\n";

    for i=0 to !num_comps - 1 do
      fprintf mout "n %d ac 0\n" (i+1);
      for j=0 to (List.length schema) - 1  do
        fprintf mout "%d " j
      done;
      fprintf mout "\n"
    done;
    fprintf mout "\n";
    close_out_noerr mout
  end;
  
  let learning_time = (t2 -. t1) in
  ignore(learning_time);
  (*nlogf "result %s %d %f %f %f %f %f \n" !network !num_comps !prior data_llg !valid_llg !test_llg learning_time; *)
  nlogf "Average data likelihood: %.5f\n" data_llg;
  flush_all()
  
;;
main ()
